import pandas as pd
from sqlalchemy import create_engine
import os

# ================== 配置 ==================
CSV_FILE = './上市公司完整信息.csv'
DB_CONFIG = {
    'host': 'localhost',
    'port': 3306,
    'user': 'taomu',
    'password': 'taomu-rcmu#2025',  # 修改为你的密码
    'database': 'taomu_stock',     # 确保数据库已创建
    'charset': 'utf8mb4'
}

# 创建数据库连接
engine = create_engine(
    f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@"
    f"{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}?"
    f"charset={DB_CONFIG['charset']}"
)

# ================== 1. 读取 CSV ==================
if not os.path.exists(CSV_FILE):
    raise FileNotFoundError(f"找不到文件: {CSV_FILE}")

df = pd.read_csv(CSV_FILE)

print("📊 原始数据预览：")
print(df.head())

# ================== 2. 中文列名 → 英文列名 ==================
column_mapping = {
    "序号": "id",
    "股票代码": "stock_code",
    "股票简称": "stock_name",
    "公司名称": "company_name",
    "省份": "province",
    "城市": "city",
    "主营业务收入(202503)": "revenue_2025q1",
    "净利润(202503)": "net_profit_2025q1",
    "员工人数": "employee_count",
    "上市日期": "listing_date",
    "招股书": "prospectus_url",
    "公司财报": "financial_report_url",
    "行业分类": "industry_category",
    "产品类型": "product_type",
    "主营业务": "main_business"
}

df_en = df.rename(columns=column_mapping)

df_en['stock_code'] = df_en['stock_code'].astype(str).str.zfill(6)

# 处理日期
df_en['listing_date'] = pd.to_datetime(df_en['listing_date'], errors='coerce')

# NaN 转为 None（MySQL 兼容）
df_en = df_en.where(pd.notnull(df_en), None)

# ================== 3. 提取财务数据 ==================
financial_data = []

for _, row in df_en.iterrows():
    stock_code = row['stock_code']

    # 提取 2025Q1 财务数据
    revenue = row.get('revenue_2025q1')
    net_profit = row.get('net_profit_2025q1')

    if pd.notna(revenue) or pd.notna(net_profit):
        financial_data.append({
            'stock_code': stock_code,
            'report_period': '2025-03-31',  # Q1 报告期
            'revenue': float(revenue) if pd.notna(revenue) else None,
            'net_profit': float(net_profit) if pd.notna(net_profit) else None
        })

df_financial = pd.DataFrame(financial_data)
df_financial['report_period'] = pd.to_datetime(df_financial['report_period'])

# ================== 4. 写入 MySQL ==================

# 表名
TABLE_COMPANIES = 'taomu_stock_listed_companies'
TABLE_FINANCIALS = 'taomu_stock_financials'

try:
    # 写入公司基本信息
    df_en[['id', 'stock_code', 'stock_name', 'company_name', 'province', 'city',
           'employee_count', 'listing_date', 'prospectus_url', 'financial_report_url',
           'industry_category', 'product_type', 'main_business']] \
      .to_sql(
        name=TABLE_COMPANIES,
        con=engine,
        if_exists='replace',  # 第一次用 replace，后续建议 append 或 update
        index=False,
        method='multi',
        chunksize=1000
    )
    print("✅ 公司基本信息已写入 MySQL 表 `listed_companies`")

    # 写入财务数据
    df_financial.to_sql(
        name=TABLE_FINANCIALS,
        con=engine,
        if_exists='append',  # 财务数据建议追加
        index=False,
        method='multi',
        chunksize=1000
    )
    print("✅ 财务数据已写入 MySQL 表 `financials`")

except Exception as e:
    print(f"❌ 写入失败: {e}")
